{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Benchmarking Different Versions of a Small Language Model Before Deployment on an Endpoint.\n",
        "This notebook is a companion of chapter 10 of the \"Domain Specific LLMs in Action\" book, author Guglielmo Iozzia, [Manning Publications](https://www.manning.com/), 2024.  \n",
        "The code in this notebook shows how to benchmark different versions of the [GPT-2 small](https://huggingface.co/openai-community/gpt2) model to assess which one would be the most performant and the final candidate for deployment on a FastAPI endpoint. The same code applies to any other Open Source LLM hosted in the HF's Hub by replacing the model id. No hardware acceleration is needed for this model. Depending on the model under benchmark a GPU would be required.\n",
        "More details about the code can be found in the related book's chapter."
      ],
      "metadata": {
        "id": "SzvvOAdqd1ks"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Install the missing dependencies in the Colab VM (only ONNX and the ONNX runtime)."
      ],
      "metadata": {
        "id": "IOW1p7XLfbCX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install onnxruntime onnx"
      ],
      "metadata": {
        "id": "uUkIJU2RnY9G"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Load the vanilla GPT-2 small model and related tokenizer from the HF's Hub."
      ],
      "metadata": {
        "id": "8d3b-SFffjoG"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uH3WCiwujS03"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "from transformers import GPT2Tokenizer, GPT2Model\n",
        "\n",
        "device = \"cpu\"\n",
        "model_id = 'openai-community/gpt2'\n",
        "tokenizer = GPT2Tokenizer.from_pretrained(model_id)\n",
        "model = GPT2Model.from_pretrained(model_id)\n",
        "model.eval()\n",
        "model_path = Path(\"gpt2\")\n",
        "model.save_pretrained(model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Get some model architecture values, required to convert then to ONNX."
      ],
      "metadata": {
        "id": "mHCECyBFftsH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "num_layer = model.config.n_layer\n",
        "num_attention_heads = model.config.n_head\n",
        "hidden_size = model.config.n_embd"
      ],
      "metadata": {
        "id": "nEI10Hln6789"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a very simple function to benchmark the different versions of the GPT-2 small model (latency only)."
      ],
      "metadata": {
        "id": "g8jPv4PKf-Sm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import timeit\n",
        "\n",
        "def benchmark(f, name=\"\"):\n",
        "    for _ in range(10):\n",
        "        f()\n",
        "    seconds_per_iter = timeit.timeit(f, number=100) / 100\n",
        "    print(\n",
        "        f\"{name}:\",\n",
        "        f\"{seconds_per_iter * 1000:.3f} ms\",\n",
        "    )"
      ],
      "metadata": {
        "id": "F1_Fu12Ljun5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Set a prompt an tokenize it for text generation with the vanilla model."
      ],
      "metadata": {
        "id": "Do7WfLjggIhi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "text = \"Today is Saturday and\"\n",
        "inputs_base = tokenizer(text, return_tensors=\"pt\").to(device)"
      ],
      "metadata": {
        "id": "Fh3uzY9soH-u"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Benchmark the vanilla model.\n",
        "\n"
      ],
      "metadata": {
        "id": "8TMfizPWgOCZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "benchmark(lambda: model(**inputs_base), \"PyTorch\")"
      ],
      "metadata": {
        "id": "tHFJhH1Z1bje"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Convert the vanilla model to ONNX format"
      ],
      "metadata": {
        "id": "TbFwgpojgR2K"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from transformers import BatchEncoding\n",
        "\n",
        "input_ids: BatchEncoding = tokenizer(\n",
        "    text, add_special_tokens=True, return_attention_mask=False, return_tensors=\"pt\"\n",
        ")\n",
        "for k, v in input_ids.items():\n",
        "    input_ids[k] = v.type(dtype=torch.int32)\n",
        "input_tensor = input_ids['input_ids']\n",
        "onnx_model_path = \"gpt2_onnx.onnx\"\n",
        "dynamic_axes = {\n",
        "    'input_ids': {0: 'batch_size', 1: 'sequence_length'},\n",
        "    'logits': {0: 'batch_size', 1: 'sequence_length'},\n",
        "}\n",
        "\n",
        "torch.onnx.export(\n",
        "    model,\n",
        "    f=onnx_model_path,\n",
        "    args= (input_tensor,),\n",
        "    input_names=['input_ids'],\n",
        "    output_names=['logits'],\n",
        "    quantization=False,\n",
        "    var_output_seq=True,\n",
        "    do_constant_folding=True,\n",
        "    opset_version=18,\n",
        "    dynamic_axes=dynamic_axes\n",
        ")"
      ],
      "metadata": {
        "id": "szNfYJntjkbT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Just in case, if you need to focus only one version of the model and run larger benchmarks on it, and need to free memory, remove the vanilla model (not needed anymore)."
      ],
      "metadata": {
        "id": "wcdTHBtAgeBD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "del model\n",
        "\n",
        "import gc\n",
        "gc.collect()"
      ],
      "metadata": {
        "id": "1tPfDCrQxZ1_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Create an inference session for the ONNX model."
      ],
      "metadata": {
        "id": "WWflEcrlgmF_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from onnxruntime import InferenceSession\n",
        "\n",
        "providers=[\"CPUExecutionProvider\"]\n",
        "sess = InferenceSession(onnx_model_path,\n",
        "                        providers=providers)"
      ],
      "metadata": {
        "id": "76E_wSp713r_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Prepare the input to benchmark this and the other ONNX version of the model that we are going to build."
      ],
      "metadata": {
        "id": "kFl3AV0ugqCD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "\n",
        "encodings_dict = tokenizer.batch_encode_plus([text])\n",
        "input_ids = torch.tensor(encodings_dict[\"input_ids\"], dtype=torch.int32)\n",
        "empty_past = []\n",
        "batch_size = input_ids.size(0)\n",
        "sequence_length = input_ids.size(1)\n",
        "past_shape = [2, batch_size, num_attention_heads, 0, hidden_size // num_attention_heads]\n",
        "for i in range(num_layer):\n",
        "    empty_past.append(torch.empty(past_shape).type(torch.float32).to(device))\n",
        "\n",
        "ort_inputs = {\n",
        "    \"input_ids\": input_ids.cpu().numpy()\n",
        "}"
      ],
      "metadata": {
        "id": "gg2v2XtR1f8M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Benchmark the ONNX model."
      ],
      "metadata": {
        "id": "bquSzhflg79t"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "benchmark(lambda: sess.run(None, ort_inputs), \"ONNX\")"
      ],
      "metadata": {
        "id": "GyQl2YZpkqR6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "To free memory, remove the ONNX model inference session (not needed anymore)."
      ],
      "metadata": {
        "id": "3kyuqQxvhEzV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "del sess\n",
        "\n",
        "import gc\n",
        "gc.collect()"
      ],
      "metadata": {
        "id": "PcYOS5XQ1Qd3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Optimize the ONNX model."
      ],
      "metadata": {
        "id": "hx9a1EwlhJx0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from onnxruntime.transformers.optimizer import optimize_model\n",
        "\n",
        "optimized_onnx_path = \"gpt2_optimized.onnx\"\n",
        "optimized_model = optimize_model(input=onnx_model_path,\n",
        "                                 model_type=\"gpt2\",\n",
        "                                 use_gpu=False)\n",
        "optimized_model.save_model_to_file(optimized_onnx_path)"
      ],
      "metadata": {
        "id": "XFFVX-JhjoWm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Create an inference session for the optimized ONNX model."
      ],
      "metadata": {
        "id": "GeopIxKUhNOJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "optimized_sess = InferenceSession(optimized_onnx_path,\n",
        "                                  providers=providers)"
      ],
      "metadata": {
        "id": "mCq0BcUH2A0B"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Benchmark the optimized ONNX model."
      ],
      "metadata": {
        "id": "76_T_CkqhR9w"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "benchmark(lambda: optimized_sess.run(None, input_feed=ort_inputs),\n",
        "          \"ONNX optimized\")"
      ],
      "metadata": {
        "id": "4SL5pQPd1i8b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Just in case, if you need to focus only one version of the model and run larger benchmarks on it, and need to free memory, remove the optimize ONNX model inference session (not needed anymore)."
      ],
      "metadata": {
        "id": "wcHi6w0vhVSR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "del optimized_sess\n",
        "\n",
        "import gc\n",
        "gc.collect()"
      ],
      "metadata": {
        "id": "sU2QXa9R1jSg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Downsize the optimized ONNX model to FP16."
      ],
      "metadata": {
        "id": "o5CLqOBEhbMq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from copy import deepcopy\n",
        "\n",
        "optimized_fp16_model_path = \"optimized_fp16.onnx\"\n",
        "optimized_fp16_model = deepcopy(optimized_model)\n",
        "optimized_fp16_model.convert_float_to_float16()\n",
        "optimized_fp16_model.save_model_to_file(optimized_fp16_model_path)"
      ],
      "metadata": {
        "id": "3Z12zsZTjrXT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "To free memory, remove the optimized ONNX model (not needed anymore)."
      ],
      "metadata": {
        "id": "mrmNiTX1hiRy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "del optimized_model\n",
        "\n",
        "import gc\n",
        "gc.collect()"
      ],
      "metadata": {
        "id": "IAh82sUh2Twf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Create an inference session for the optimized FP16 ONNX model."
      ],
      "metadata": {
        "id": "4ra89dEmhk-b"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "optimized_fp16_sess = InferenceSession(\n",
        "    optimized_fp16_model_path, providers=providers\n",
        ")"
      ],
      "metadata": {
        "id": "dpTp37CsoosO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Benchmark the optimized FP16 ONNX model."
      ],
      "metadata": {
        "id": "lg_Fp-YjhrcP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "benchmark(lambda: optimized_fp16_sess.run(None, input_feed=ort_inputs),\n",
        "          \"ONNX optimized fp16\")"
      ],
      "metadata": {
        "id": "Q7YWc2ZAo3XQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the benchmark for each version of the model using different values of `max_length` (1, 4, 64, 256, 512, 1024) and then compare the results."
      ],
      "metadata": {
        "id": "1nzlrarHhyKK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np"
      ],
      "metadata": {
        "id": "-NVM-N76SN-O"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "MAX_SEQUENCE_LENGTH=1024\n",
        "for n in [1, 4, 64, 256, 512, 1024]:\n",
        "    print(f\"====== Tokens {n} ======\")\n",
        "    txt = \" \".join([\"word\"] * n)\n",
        "\n",
        "    pt_inputs_base = tokenizer(txt,\n",
        "                              max_length=MAX_SEQUENCE_LENGTH, return_tensors=\"pt\").to(device)\n",
        "    pt_inputs = tokenizer(txt,\n",
        "                              max_length=MAX_SEQUENCE_LENGTH, return_tensors=\"pt\").to(device)\n",
        "    ort_inputs = dict(tokenizer(txt,\n",
        "                              max_length=MAX_SEQUENCE_LENGTH,\n",
        "                              return_tensors=\"np\",\n",
        "                              return_attention_mask=False))\n",
        "    ort_inputs['input_ids'] = ort_inputs['input_ids'].astype(np.int32)\n",
        "\n",
        "    benchmark(lambda: model(**pt_inputs), f\"Pytorch ({n} tokens)\")\n",
        "    benchmark(lambda: sess.run(None, {'input_ids': ort_inputs['input_ids']}), f\"ONNX ({n} tokens)\")\n",
        "    benchmark(lambda: optimized_sess.run(None, ort_inputs), f\"ONNX optimized ({n} tokens)\")\n",
        "    benchmark(\n",
        "        lambda: optimized_fp16_sess.run(None, ort_inputs),\n",
        "        f\"ONNX optimized fp16 ({n} tokens)\",\n",
        "    )"
      ],
      "metadata": {
        "id": "mCEN-PT_p0aT"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}